Rethinking Learning Rate and Batch Size (Part 4): EMA Effects

In "Rethinking Learning Rate and Batch Size (Part 2): Mean Field Theory", we mentioned that one reason for focusing on SignSGD is that we typically use it as a theoretical approximation for Adam—a common simplification strategy when analyzing Adam theoretically. Beyond learning rate analysis, we have also employed this simplification in contexts such as "Configuring Different Learning Rates: Can LoRA Improve Further?" and "Exploring MuP: Cross-Model Scale Transfer of Hyperparameters".

However, is SignSGD truly a good approximation for Adam? An obvious discrepancy is that SignSGD's Update RMS is always 1, whereas Adam's is not. I discovered that the core reason for this difference is momentum, which is ubiquitous in optimizers like Adam, Lion, and Muon. Therefore, in this article, we examine the effects of momentum—more broadly, Exponential Moving Averages (EMA).

Problem Analysis#

From the perspective of Adam, SignSGD corresponds to the special case where $\beta_1=\beta_2=0$, or corresponds to Adam's first update step (regardless of $\beta_1,\beta_2$). Therefore, we believe it must share some commonalities with Adam and can capture some general patterns.

This leads to the following analytical process. In summary, we can confirm that the role of $\epsilon$ is indeed secondary; the true protagonist is actually momentum—which is essentially a "moving average" of gradients—and this is precisely the main focus of this article: "EMA (Exponential Moving Average)".

Gradient Descent#

To analyze the variations introduced by EMA, we start with SGDM (SGD with momentum). In practice, when using SGD, it is rare not to include momentum:

In practical use, $\boldsymbol{g}_t$ is replaced by $\tilde{\boldsymbol{g}}_{B,t}$, a random variable with mean $\boldsymbol{g}_t$ and covariance matrix $\boldsymbol{\Sigma}_t/B$. These basic settings are the same as in "Rethinking Learning Rate and Batch Size (Part 1): Current Landscape". The noise here arises from randomly sampling different batches, so we can reasonably assume that $\tilde{\boldsymbol{g}}_{B,t}$ across different $t$ are mutually independent.

Our task is to compute:

The related derivations have been provided in previous articles and won't be repeated here. For SGDM, $\tilde{\boldsymbol{\varphi}}_B = \boldsymbol{m}_t$, which can be expanded as:

Batch Size Amplification#

Now we can compute:

We further assume that when model training enters a "steady state," the gradient changes slowly. Thus, we can approximate $\boldsymbol{g}_s$ with the current gradient $\boldsymbol{g}_t$, yielding:

As for $\mathbb{E}[\boldsymbol{m}_t \boldsymbol{m}_t^{\top}]$, we use the identity $\mathbb{E}[\boldsymbol{m}_t \boldsymbol{m}_t^{\top}] = \mathbb{E}[\boldsymbol{m}_t] \mathbb{E}[\boldsymbol{m}_t]^{\top} + \mathbb{C}\text{ov}[\boldsymbol{m}_t,\boldsymbol{m}_t]$, then utilize the additivity of variance:

Similarly, assuming slow variation of the covariance matrix, we have:

Substituting into equation (2) yields:

Signed Momentum#

Further, we consider SignSGDM, which can be viewed as a special case of Lion—essentially SGDM with an added $\newcommand{sign}{\mathop{\text{sign}}}\sign$:

In actual training, $\boldsymbol{g}_t$ is similarly replaced by $\tilde{\boldsymbol{g}}_{B,t}$. For SignSGDM, $\tilde{\boldsymbol{\varphi}}_B = \sign(\boldsymbol{m}_t)$. Using the mean field approximation:

Here, vector multiplication defaults to the Hadamard product. The numerator $\mathbb{E}[\boldsymbol{m}_t]$ was computed in the previous section. The denominator $\mathbb{E}[\boldsymbol{m}_t^2]$ actually equals $\newcommand{diag}{\mathop{\text{diag}}}\diag(\mathbb{E}[\boldsymbol{m}_t \boldsymbol{m}_t^{\top}])$, so we can substitute the results from the previous section to obtain:

where $\boldsymbol{\sigma}_t^2 = \diag(\boldsymbol{\Sigma}_t)$, and $\mathcal{B}_{\text{simple}} = \tr(\boldsymbol{\Sigma}_t)/\boldsymbol{g}_t^{\top}\boldsymbol{g}_t$. The above equation is equivalent to SignSGD with $B$ replaced by $\frac{1 + \beta_1}{1 - \beta_1}B$. If we further compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B\tilde{\boldsymbol{\varphi}}_B^{\top}]$, we would find the same conclusion. Thus, similar to SGDM, momentum effectively amplifies SignSGD's batch size by a factor of $\frac{1 + \beta_1}{1 - \beta_1}$.

In "Rethinking Learning Rate and Batch Size (Part 3): Muon", we calculated Muon's learning rate scaling pattern and found it to be consistent with SignSGD. Therefore, we can assert that the role of momentum in Muon is similar to that in SignSGDM—both approximately amplify the batch size by a factor of $\frac{1 + \beta_1}{1 - \beta_1}$.

Double Smoothing#

Finally, we examine Adam:

In actual training, $\boldsymbol{g}_t$ is replaced by $\tilde{\boldsymbol{g}}_{B,t}$. We consider the state where training has already entered a "steady state," i.e., $t\to\infty$, so we do not distinguish between $\boldsymbol{m}_t$ and $\hat{\boldsymbol{m}}_t$, or between $\boldsymbol{v}_t$ and $\hat{\boldsymbol{v}}_t$. Additionally, to focus on the effect of EMA, we set $\epsilon = 0$. Then for Adam, $\tilde{\boldsymbol{\varphi}}_B=\boldsymbol{m}_t/\sqrt{\boldsymbol{v}_t}$. The difference from SignSGDM is that the denominator $\boldsymbol{m}_t^2$ is replaced by another EMA statistic $\boldsymbol{v}_t$.

Using the mean field approximation:

We have already computed $\mathbb{E}[\boldsymbol{m}_t]$, so we only need to compute $\mathbb{E}[\boldsymbol{v}_t]$:

As before, the last approximation assumes slow variation of gradients and variances, and $t\to\infty$. Thus, we have:

This result is indeed identical to SignSGD, so from the perspective of the first moment, SignSGD as an approximation for Adam is reasonable. However, we also have the second moment $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B \tilde{\boldsymbol{\varphi}}_B^{\top}]$. Under the assumption of component independence, we only need to compute $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B^2]$:

Two Special Cases#

We examine two special cases. First, when $\beta_1=0$, the numerator and denominator are identical, making $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B^2]$ a vector of all ones—consistent with SignSGD. Therefore, SignSGD is a good approximation for Adam with $\beta_1=0$—which is RMSProp. As $\beta_1$ increases, the approximation quality begins to degrade.

When $\beta_1=1$, we have:

This yields $\mathbb{E}[\tilde{\boldsymbol{\varphi}}_B \tilde{\boldsymbol{\varphi}}_B^{\top}] \approx \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B] \mathbb{E}[\tilde{\boldsymbol{\varphi}}_B]^{\top}$. Substituting into equation (2):

General Analysis#

Let's rewrite equation (16):

From this, we can write:

Then:

Here, the subscriptless $\beta$ equals $(1 + \mathcal{B}_{\text{simple}}/B)^{-1/2}$. Without careful reading, this might be confused with $\beta_1,\beta_2$—I apologize, as this is notation from the previous two articles and we continue using it here. Unlike SignSGD, where assuming a diagonal Hessian matrix eliminates the Surge phenomenon, the above expression still exhibits the Surge phenomenon even under a diagonal Hessian assumption. In that case:

By the inequality of arithmetic and geometric means, the above expression reaches its maximum at $\beta^*=\sqrt{\frac{1-\beta_1}{2\beta_1}}$. However, note that by definition $\beta\in(0,1)$, so we must also check whether $\beta^*\in(0,1)$, i.e., $\beta_1 > 1/3$. If this condition is not satisfied, the maximum still occurs at $\beta=1$, and there is no Surge phenomenon. Conversely, when $\beta_1 > 1/3$ and $\beta > \beta^*$ (i.e., $B > \frac{1-\beta_1}{3\beta_1-1}\mathcal{B}_{\text{simple}}$), the learning rate should decrease as batch size increases.

Summary#

This article provides a preliminary analysis of how the EMA mechanism in optimizers affects learning rate and batch size scaling laws. It confirms that EMA, particularly the momentum mechanism, subtly alters scaling laws. Optimizers like Adam, which involve dual EMA operations, exhibit some new characteristics distinct from SignSGD.